{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "2.0.0-beta0\n",
      "2.2.4-tf\n"
     ]
    }
   ],
   "source": [
    "import tensorflow as tf\n",
    "import numpy as np\n",
    "print(tf.__version__)\n",
    "print(tf.keras.__version__)\n",
    "gpus = tf.config.experimental.list_physical_devices('GPU')\n",
    "tf.config.experimental.set_memory_growth(gpus[0], True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "def conv_block(func, bottom, filters, kernel_size, strides=1, dilation_rate=1, reg=1e-4, apply_bn=True, apply_relu=True, training=False):\n",
    "    conv_params = {\n",
    "        'padding': 'same',\n",
    "        'dilation_rate': dilation_rate,\n",
    "        'kernel_initializer': tf.keras.initializers.glorot_normal,\n",
    "        'kernel_regularizer': tf.keras.regularizers.l2(reg),\n",
    "        'bias_regularizer': tf.keras.regularizers.l2(reg)}\n",
    "    bottom = func(filters, kernel_size, strides, **conv_params)(bottom)\n",
    "    if apply_bn:\n",
    "        bottom = tf.keras.layers.BatchNormalization()(bottom, training=training)\n",
    "    if apply_relu:\n",
    "        bottom = tf.keras.layers.ReLU()(bottom)\n",
    "    return bottom"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "def res_block(func, bottom, filters, kernel_size, strides=1, dilation_rate=1, reg=1e-4, training=False):\n",
    "    short_cut = bottom\n",
    "    bottom = conv_block(func, bottom, filters, kernel_size, strides, dilation_rate, reg,\n",
    "                        training=training)\n",
    "    bottom = conv_block(func, bottom, filters, kernel_size, 1, 1, reg,\n",
    "                        apply_relu=False, training=training)\n",
    "    if strides != 1 or dilation_rate != 1:\n",
    "        short_cut = conv_block(func, short_cut, filters, kernel_size, strides, reg=reg,\n",
    "                               apply_bn=False, apply_relu=False, training=training)\n",
    "    bottom = tf.add(bottom, short_cut)\n",
    "    bottom = tf.keras.layers.ReLU()(bottom)\n",
    "    return bottom"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "#3d的代码实现还未验证\n",
    "def SPP_branch(func, bottom, pool_size, filters, kernel_size, strides=1, dilation_rate=1, reg=1e-4,\n",
    "               apply_bn=True, apply_relu=True, training=False):\n",
    "    if func != tf.keras.layers.Conv2D:\n",
    "        size = tf.Tensor.get_shape(bottom)\n",
    "        size = tf.TensorShape(size).as_list()\n",
    "        size_1 = size #这里是否需要copy代替\n",
    "        size_1[1] = size_1[0] * size_1[1]\n",
    "        bottom = tf.keras.layers.AveragePooling3D(pool_size, padding='same')(bottom)\n",
    "        bottom = conv_block(func, bottom, filters, kernel_size, strides, dilation_rate, reg, training=training)\n",
    "        bottom = tf.reshape(bottom, shape=size_1[1:])\n",
    "        bottom = tf.image.resize(bottom, size_1[-3:-1])\n",
    "        bottom = tf.reshape(bottom, shape=size)\n",
    "    else:\n",
    "        size = tf.Tensor.get_shape(bottom)\n",
    "        size = tf.TensorShape(size).as_list()[1:3]\n",
    "        bottom = tf.keras.layers.AveragePooling2D(pool_size, padding='same')(bottom)\n",
    "        bottom = conv_block(func, bottom, filters, kernel_size, strides, dilation_rate, reg, training=training)\n",
    "        bottom = tf.image.resize(bottom, size)\n",
    "    return bottom"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "def hourglass_stack(strs, bottom, filters_list, kernel_size_list, short_cut_list, reg=1e-4, training=False):\n",
    "    output = []\n",
    "    conv_func, deconv_func = (tf.keras.layers.Conv2D, tf.keras.layers.Conv2DTranspose) if strs == '2d' else (tf.keras.layers.Conv3D, tf.keras.layers.Conv3DTranspose)\n",
    "    for i, (filters, kernel_size, short_cut) in enumerate(zip(filters_list, kernel_size_list, short_cut_list)):\n",
    "        if i < 2:\n",
    "            bottom = conv_block(conv_func, bottom, filters, kernel_size, strides=2, reg=reg, training=training)\n",
    "            bottom = conv_block(conv_func, bottom, filters, kernel_size, reg=reg, apply_relu=False, training=training)\n",
    "            if short_cut is not None:\n",
    "                if type(short_cut) is int:\n",
    "                    short_cut = output[short_cut]\n",
    "                bottom = tf.add(bottom, short_cut)\n",
    "            bottom = tf.keras.layers.ReLU()(bottom)\n",
    "        else:\n",
    "            bottom = conv_block(deconv_func, bottom, filters, kernel_size, strides=2, reg=reg, apply_relu=False, training=training)\n",
    "            if short_cut is not None:\n",
    "                if type(short_cut) is int:\n",
    "                    short_cut = output[short_cut]\n",
    "                bottom = tf.add(bottom, short_cut)\n",
    "            if i == 2:\n",
    "                bottom = tf.keras.layers.ReLU()(bottom)\n",
    "        output.append(bottom)\n",
    "    return output"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "def soft_arg_min(filtered_cost_volume, name):\n",
    "    probability_volume = tf.nn.softmax(tf.scalar_mul(-1, filtered_cost_volume), axis=1, name='prob_volume')\n",
    "    volume_shape = tf.shape(probability_volume)\n",
    "    soft_1d = tf.cast(tf.range(0, volume_shape[1], dtype=tf.int32), dtype=tf.float32)\n",
    "    soft_4d = tf.tile(soft_1d, tf.stack([volume_shape[0] * volume_shape[2] * volume_shape[3]]))\n",
    "    soft_4d = tf.reshape(soft_4d, [volume_shape[0], volume_shape[2], volume_shape[3], volume_shape[1]])\n",
    "    soft_4d = tf.transpose(soft_4d, [0, 3, 1, 2])\n",
    "    estimated_disp_image = tf.reduce_sum(soft_4d * probability_volume, axis=1)\n",
    "    return estimated_disp_image"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "class PSMNET(tf.keras.Model):\n",
    "    def __init__(self, max_disp=192, training=False):\n",
    "        super(PSMNET, self).__init__(name='my_model')\n",
    "        self.max_disp = max_disp\n",
    "        self.reg = 1e-4\n",
    "        self.training = training\n",
    "        \n",
    "    def call(self, left, right):\n",
    "        conv4_left = self.CNN(left)\n",
    "        fusion_left = self.SPP_modile(conv4_left)\n",
    "        \n",
    "        conv4_right = self.CNN(right)\n",
    "        fusion_right = self.SPP_modile(conv4_right)\n",
    "        \n",
    "        cost_voloum = self.Cost_volume(fusion_left, fusion_right)\n",
    "        outputs = self.Hourglass(cost_voloum)\n",
    "        \n",
    "        return outputs\n",
    "        \n",
    "    def CNN(self, bottom):\n",
    "        outputs = []\n",
    "        bottom = conv_block(tf.keras.layers.Conv2D, bottom, 32, 3, strides=2, reg=self.reg, training=self.training)\n",
    "        for i in range(1, 3):\n",
    "            bottom = conv_block(tf.keras.layers.Conv2D, bottom, 32, 3, reg=self.reg, training=self.training)\n",
    "        for i in range(3):\n",
    "            bottom = res_block(tf.keras.layers.Conv2D, bottom, 32, 3, reg=self.reg, training=self.training)\n",
    "        bottom = res_block(tf.keras.layers.Conv2D, bottom, 64, 3, strides=2, reg=self.reg, training=self.training)\n",
    "        for i in range(1, 16):\n",
    "            bottom = res_block(tf.keras.layers.Conv2D, bottom, 64, 3, reg=self.reg, training=self.training)\n",
    "        outputs.append(bottom) #add conv2_16\n",
    "        for i in range(3):\n",
    "            bottom = res_block(tf.keras.layers.Conv2D, bottom, 128, 3, dilation_rate=2, reg=self.reg, training=self.training)\n",
    "        for i in range(3):\n",
    "            bottom = res_block(tf.keras.layers.Conv2D, bottom, 128, 3, dilation_rate=4, reg=self.reg, training=self.training)\n",
    "        outputs.append(bottom) #add conv4_3\n",
    "        return outputs\n",
    "    \n",
    "    def SPP_modile(self, inputs):\n",
    "        branches = []\n",
    "        for i, pool_size in enumerate([64, 32, 16, 8]):\n",
    "            branches.append(SPP_branch(tf.keras.layers.Conv2D, inputs[1], pool_size, 32, 3, reg=self.reg, training=self.training))\n",
    "        bottom = tf.concat([inputs[0], inputs[1]] + branches, axis=-1)\n",
    "        bottom = conv_block(tf.keras.layers.Conv2D, bottom, 128, 3, reg=self.reg, training=self.training)\n",
    "        bottom = conv_block(tf.keras.layers.Conv2D, bottom, 32, 1, reg=self.reg, training=self.training)\n",
    "        return bottom\n",
    "    \n",
    "    def Cost_volume(self, left, right):\n",
    "        disparity_costs = []\n",
    "        for i in range(self.max_disp // 4):\n",
    "            if i > 0:\n",
    "                left_tensor = tf.keras.backend.spatial_2d_padding(left[:, :, i:, :], padding=((0, 0), (i, 0)))\n",
    "                right_tensor = tf.keras.backend.spatial_2d_padding(right[:, :, :-i, :], padding=((0, 0), (i, 0)))\n",
    "                cost_branch = tf.concat([left_tensor, right_tensor], axis=3)\n",
    "            else:\n",
    "                cost_branch = tf.concat([left, right], axis=3)\n",
    "            disparity_costs.append(cost_branch)\n",
    "        cost_volume = tf.stack(disparity_costs, axis=1)\n",
    "        return cost_volume\n",
    "    \n",
    "    def Hourglass(self, bottom):\n",
    "        for i in range(2):\n",
    "            bottom = conv_block(tf.keras.layers.Conv3D, bottom, 32, 3, reg=self.reg, training=self.training)\n",
    "        _3Dconv1 = res_block(tf.keras.layers.Conv3D, bottom, 32, 3, reg=self.reg, training=self.training)\n",
    "        _3Dstack = [hourglass_stack('3d', _3Dconv1, [64,64,64,32], [3,3,3,3], [None,None,0,_3Dconv1], reg=self.reg, training=self.training)]\n",
    "        for i in range(1, 3):\n",
    "            _3Dstack.append(hourglass_stack('3d', _3Dstack[-1][-1], [64,64,64,32], [3,3,3,3], [_3Dstack[-1][2],None,_3Dstack[0][0],_3Dconv1], reg=self.reg, training=self.training))\n",
    "        output_1 = conv_block(tf.keras.layers.Conv3D, _3Dstack[-1][-1], 32, 3, reg=self.reg, training=self.training)\n",
    "        output_1 = conv_block(tf.keras.layers.Conv3D, output_1, 1, 3, reg=self.reg, training=self.training)\n",
    "        outputs = [output_1]\n",
    "        for i in range(1, 3):\n",
    "            output = conv_block(tf.keras.layers.Conv3D, _3Dstack[i][-1], 32, 3, reg=self.reg, training=self.training)\n",
    "            output = conv_block(tf.keras.layers.Conv3D, output, 1, 3, reg=self.reg, training=self.training)\n",
    "            outputs.append(tf.add(output, outputs[-1]))\n",
    "        return outputs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "def output(outputs):\n",
    "    disps = []\n",
    "    for i, output in enumerate(outputs):\n",
    "        squeeze = tf.squeeze(output, [4])\n",
    "        transpose = tf.transpose(squeeze, [0, 2, 3, 1])\n",
    "        upsample = tf.transpose(tf.image.resize(transpose, [self.height, self.weight]), [0, 3, 1, 2])\n",
    "        disps.append(soft_arg_min(upsample, 'soft_arg_min_%d' % (i+1)))\n",
    "    return disps\n",
    "\n",
    "def smooth_l1_loss(disps_pred, disps_targets, sigma=1.0):\n",
    "    sigma_2 = sigma ** 2\n",
    "    disps_diff = disps_pred - disps_targets\n",
    "    abs_disps_diff = tf.abs(disps_diff)\n",
    "    loss_disps = tf.where(tf.less(abs_disps_diff, 1./sigma_2), tf.pow(disps_diff, 2) * (sigma_2/2.), (abs_disps_diff - (0.5/sigma_2)))\n",
    "    loss = tf.reduce_mean(tf.reduce_mean(loss_disps, [1,2]))\n",
    "    return loss\n",
    "\n",
    "def loss(self):\n",
    "    self.disps = self.output(outputs)\n",
    "    self.loss = tf.add_n([0.5 * self.smooth_l1_loss(self.disps[0], label),\n",
    "                          0.7 * self.smooth_l1_loss(self.disps[1], label),\n",
    "                          1.0 * self.smooth_l1_loss(self.disps[2], label)])\n",
    "\n",
    "def train(self, left, right, label, lr=0.001):\n",
    "    with tf.GradientTape() as tape:\n",
    "        loss = self.loss\n",
    "        optimizer = tf.keras.optimizers.Adam(learning_rate=lr)\n",
    "        gradients = tape.gradient(loss, model.trainable_variables)\n",
    "        optimizer.apply_gradients(zip(gradients, model.trainable_variables))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "WARNING:tensorflow:Entity <bound method PSMNET.call of <__main__.PSMNET object at 0x000001E5BBE5D5F8>> could not be transformed and will be executed as-is. Please report this to the AutoGraph team. When filing the bug, set the verbosity to 10 (on Linux, `export AUTOGRAPH_VERBOSITY=10`) and attach the full output. Cause: converting <bound method PSMNET.call of <__main__.PSMNET object at 0x000001E5BBE5D5F8>>: AssertionError: Bad argument number for Name: 3, expecting 4\n",
      "WARNING: Entity <bound method PSMNET.call of <__main__.PSMNET object at 0x000001E5BBE5D5F8>> could not be transformed and will be executed as-is. Please report this to the AutoGraph team. When filing the bug, set the verbosity to 10 (on Linux, `export AUTOGRAPH_VERBOSITY=10`) and attach the full output. Cause: converting <bound method PSMNET.call of <__main__.PSMNET object at 0x000001E5BBE5D5F8>>: AssertionError: Bad argument number for Name: 3, expecting 4\n"
     ]
    }
   ],
   "source": [
    "model = PSMNET()\n",
    "left = tf.keras.Input([256, 512 ,3], name='left')\n",
    "right = tf.keras.Input([256, 512 ,3], name='right')\n",
    "output = model(left, right)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "a = tf.cast(np.random.ranf([4,256,512,3]), dtype=tf.float32)\n",
    "b = tf.cast(np.random.ranf([4,256,512,3]), dtype=tf.float32)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "x = model(a,b)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "y = model(a,b)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "xa = np.asarray(x)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ya = np.array(y)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "psmnet = PSMNET(max_disp=192, training=True)\n",
    "inputs_label = tf.keras.Input((256, 512), name='left_lable')\n",
    "inputs1 = tf.keras.Input((256, 512 ,3), name='left')\n",
    "bottom1 = psmnet.CNN(inputs1)\n",
    "bottom1 = psmnet.SPP_modile(bottom1)\n",
    "\n",
    "inputs2 = tf.keras.Input((256, 512 ,3), name='right')\n",
    "bottom2 = psmnet.CNN(inputs2)\n",
    "bottom2 = psmnet.SPP_modile(bottom2)\n",
    "\n",
    "bottom = psmnet.Cost_volume(bottom1, bottom2)\n",
    "outputs = psmnet.Hourglass(bottom)\n",
    "\n",
    "disps = psmnet.output(outputs)\n",
    "loss = 0.5*psmnet.smooth_l1_loss(disps[0], inputs_label) + 0.7*psmnet.smooth_l1_loss(disps[1], inputs_label) + psmnet.smooth_l1_loss(disps[2], inputs_label)\n",
    "\n",
    "model = tf.keras.Model(inputs=[inputs1, inputs2], outputs=outputs, name='psmnet model')\n",
    "model.summary(line_length=120)\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "(train_images, train_labels), (test_images, test_labels) = tf.keras.datasets.mnist.load_data('mnist.npz')\n",
    "# 6万张训练图片，1万张测试图片\n",
    "train_images = train_images.reshape((60000, 28, 28, 1))\n",
    "test_images = test_images.reshape((10000, 28, 28, 1))\n",
    "# 像素值映射到 0 - 1 之间\n",
    "train_images, test_images = train_images / 255.0, test_images / 255.0"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "inputs = tf.keras.Input((28,28,1), name='mnist')\n",
    "bottom = tf.keras.layers.Conv2D(32, 3, activation='relu')(inputs)\n",
    "bottom = tf.keras.layers.MaxPooling2D(2)(bottom)\n",
    "bottom = tf.keras.layers.Conv2D(64, 3, activation='relu')(bottom)\n",
    "bottom = tf.keras.layers.MaxPooling2D(2)(bottom)\n",
    "bottom = tf.keras.layers.Conv2D(64, (3, 3), activation='relu')(bottom)\n",
    "bottom = tf.keras.layers.Flatten()(bottom)\n",
    "bottom = tf.keras.layers.Dense(64, activation='relu')(bottom)\n",
    "outputs = tf.keras.layers.Dense(10, activation='softmax')(bottom)\n",
    "\n",
    "model = tf.keras.Model(inputs=inputs, outputs=outputs, name='mnist model')\n",
    "model.summary()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def train():\n",
    "    check_path = './ckpt/cp-{epoch:04d}.ckpt'\n",
    "    # period 每隔5epoch保存一次\n",
    "    save_model_cb = tf.keras.callbacks.ModelCheckpoint(check_path, save_weights_only=True, verbose=1, period=5)\n",
    "    model.compile(optimizer='adam',\n",
    "                  loss='sparse_categorical_crossentropy',\n",
    "                  metrics=['accuracy'])\n",
    "    model.fit(train_images,\n",
    "               train_labels,\n",
    "               epochs=5,\n",
    "               callbacks=[save_model_cb])\n",
    "def test():\n",
    "    latest = tf.train.latest_checkpoint('./ckpt')\n",
    "    model.compile(optimizer='adam',\n",
    "                  loss='sparse_categorical_crossentropy',\n",
    "                  metrics=['accuracy'])\n",
    "    model.load_weights(latest)\n",
    "    test_loss, test_acc = model.evaluate(test_images, test_labels)\n",
    "    print(\"准确率: %.4f，共测试了%d张图片 \" % (test_acc, len(test_labels)))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "train()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "test()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "strs = '2d'\n",
    "conv_func, deconv_func = (tf.keras.layers.Conv2D, tf.keras.layers.Conv2DTranspose) if strs == '2d' else (tf.keras.layers.Conv3D, tf.keras.layers.Conv3DTranspose)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "conv_func"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for i in range(1,3):\n",
    "    print(i)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class PSMNET():\n",
    "    def __init__(self, height=256, weight=512, max_disp=192, training=False):\n",
    "        #super(PSMNET, self).__init__(name='PSMNET')\n",
    "        self.height = height\n",
    "        self.weight = weight\n",
    "        self.max_disp = max_disp\n",
    "        self.reg = 1e-4\n",
    "        self.training = training\n",
    "        self.build_model()\n",
    "        \n",
    "    def build_model(self):\n",
    "        left = tf.keras.Input([256, 512 ,3], name='left')\n",
    "        right = tf.keras.Input([256, 512 ,3], name='right')\n",
    "        \n",
    "        conv4_left = self.CNN(left)\n",
    "        fusion_left = self.SPP_modile(conv4_left)\n",
    "        \n",
    "        conv4_right = self.CNN(right)\n",
    "        fusion_right = self.SPP_modile(conv4_right)\n",
    "        \n",
    "        cost_voloum = self.Cost_volume(fusion_left, fusion_right)\n",
    "        outputs = self.Hourglass(cost_voloum)\n",
    "        \n",
    "        #model.summary(line_length=120)\n",
    "        \n",
    "        self.disps = self.output(outputs)\n",
    "        self.loss = tf.add_n([0.5 * self.smooth_l1_loss(self.disps[0], label),\n",
    "                              0.7 * self.smooth_l1_loss(self.disps[1], label),\n",
    "                              1.0 * self.smooth_l1_loss(self.disps[2], label)])\n",
    "    \n",
    "    @tf.function\n",
    "    def train(self, left, right, label, lr=0.001):\n",
    "        with tf.GradientTape() as tape:\n",
    "            loss = self.loss\n",
    "            optimizer = tf.keras.optimizers.Adam(learning_rate=lr)\n",
    "            gradients = tape.gradient(loss, model.trainable_variables)\n",
    "            optimizer.apply_gradients(zip(gradients, model.trainable_variables))\n",
    "        \n",
    "        \n",
    "    def CNN(self, bottom):\n",
    "        outputs = []\n",
    "        bottom = conv_block(tf.keras.layers.Conv2D, bottom, 32, 3, strides=2, reg=self.reg, training=self.training)\n",
    "        for i in range(1, 3):\n",
    "            bottom = conv_block(tf.keras.layers.Conv2D, bottom, 32, 3, reg=self.reg, training=self.training)\n",
    "        for i in range(3):\n",
    "            bottom = res_block(tf.keras.layers.Conv2D, bottom, 32, 3, reg=self.reg, training=self.training)\n",
    "        bottom = res_block(tf.keras.layers.Conv2D, bottom, 64, 3, strides=2, reg=self.reg, training=self.training)\n",
    "        for i in range(1, 16):\n",
    "            bottom = res_block(tf.keras.layers.Conv2D, bottom, 64, 3, reg=self.reg, training=self.training)\n",
    "        outputs.append(bottom) #add conv2_16\n",
    "        for i in range(3):\n",
    "            bottom = res_block(tf.keras.layers.Conv2D, bottom, 128, 3, dilation_rate=2, reg=self.reg, training=self.training)\n",
    "        for i in range(3):\n",
    "            bottom = res_block(tf.keras.layers.Conv2D, bottom, 128, 3, dilation_rate=4, reg=self.reg, training=self.training)\n",
    "        outputs.append(bottom) #add conv4_3\n",
    "        return outputs\n",
    "    \n",
    "    def SPP_modile(self, inputs):\n",
    "        branches = []\n",
    "        for i, pool_size in enumerate([64, 32, 16, 8]):\n",
    "            branches.append(SPP_branch(tf.keras.layers.Conv2D, inputs[1], pool_size, 32, 3, reg=self.reg, training=self.training))\n",
    "        bottom = tf.concat([inputs[0], inputs[1]] + branches, axis=-1)\n",
    "        bottom = conv_block(tf.keras.layers.Conv2D, bottom, 128, 3, reg=self.reg, training=self.training)\n",
    "        bottom = conv_block(tf.keras.layers.Conv2D, bottom, 32, 1, reg=self.reg, training=self.training)\n",
    "        return bottom\n",
    "    \n",
    "    def Cost_volume(self, left, right):\n",
    "        disparity_costs = []\n",
    "        for i in range(self.max_disp // 4):\n",
    "            if i > 0:\n",
    "                left_tensor = tf.keras.backend.spatial_2d_padding(left[:, :, i:, :], padding=((0, 0), (i, 0)))\n",
    "                right_tensor = tf.keras.backend.spatial_2d_padding(right[:, :, :-i, :], padding=((0, 0), (i, 0)))\n",
    "                cost_branch = tf.concat([left_tensor, right_tensor], axis=3)\n",
    "            else:\n",
    "                cost_branch = tf.concat([left, right], axis=3)\n",
    "            disparity_costs.append(cost_branch)\n",
    "        cost_volume = tf.stack(disparity_costs, axis=1)\n",
    "        return cost_volume\n",
    "    \n",
    "    def Hourglass(self, bottom):\n",
    "        for i in range(2):\n",
    "            bottom = conv_block(tf.keras.layers.Conv3D, bottom, 32, 3, reg=self.reg, training=self.training)\n",
    "        _3Dconv1 = res_block(tf.keras.layers.Conv3D, bottom, 32, 3, reg=self.reg, training=self.training)\n",
    "        _3Dstack = [hourglass_stack('3d', _3Dconv1, [64,64,64,32], [3,3,3,3], [None,None,0,_3Dconv1], reg=self.reg, training=self.training)]\n",
    "        for i in range(1, 3):\n",
    "            _3Dstack.append(hourglass_stack('3d', _3Dstack[-1][-1], [64,64,64,32], [3,3,3,3], [_3Dstack[-1][2],None,_3Dstack[0][0],_3Dconv1], reg=self.reg, training=self.training))\n",
    "        output_1 = conv_block(tf.keras.layers.Conv3D, _3Dstack[-1][-1], 32, 3, reg=self.reg, training=self.training)\n",
    "        output_1 = conv_block(tf.keras.layers.Conv3D, output_1, 1, 3, reg=self.reg, training=self.training)\n",
    "        outputs = [output_1]\n",
    "        for i in range(1, 3):\n",
    "            output = conv_block(tf.keras.layers.Conv3D, _3Dstack[i][-1], 32, 3, reg=self.reg, training=self.training)\n",
    "            output = conv_block(tf.keras.layers.Conv3D, output, 1, 3, reg=self.reg, training=self.training)\n",
    "            outputs.append(tf.add(output, outputs[-1]))\n",
    "        return outputs\n",
    "    \n",
    "    def output(self, outputs):\n",
    "        disps = []\n",
    "        for i, output in enumerate(outputs):\n",
    "            squeeze = tf.squeeze(output, [4])\n",
    "            transpose = tf.transpose(squeeze, [0, 2, 3, 1])\n",
    "            upsample = tf.transpose(tf.image.resize(transpose, [self.height, self.weight]), [0, 3, 1, 2])\n",
    "            disps.append(soft_arg_min(upsample, 'soft_arg_min_%d' % (i+1)))\n",
    "        return disps\n",
    "    \n",
    "    def smooth_l1_loss(self, disps_pred, disps_targets, sigma=1.0):\n",
    "        sigma_2 = sigma ** 2\n",
    "        disps_diff = disps_pred - disps_targets\n",
    "        abs_disps_diff = tf.abs(disps_diff)\n",
    "        loss_disps = tf.where(tf.less(abs_disps_diff, 1./sigma_2), tf.pow(disps_diff, 2) * (sigma_2/2.), (abs_disps_diff - (0.5/sigma_2)))\n",
    "        loss = tf.reduce_mean(tf.reduce_mean(loss_disps, [1,2]))\n",
    "        return loss"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
